#!/usr/bin/env python3
"""
数据集合并脚本
合并features.csv(全部数据) + maloss.csv(label=1) + robust.csv(label=1)
"""

import pandas as pd
import numpy as np
from pathlib import Path
import sys
from typing import List, Optional

def load_and_validate_csv(file_path: str) -> Optional[pd.DataFrame]:
    """
    加载并验证CSV文件
    """
    try:
        if not Path(file_path).exists():
            print(f"❌ 文件不存在: {file_path}")
            return None
            
        df = pd.read_csv(file_path)
        print(f"✅ 成功加载 {file_path}")
        print(f"   数据形状: {df.shape}")
        print(f"   列名: {list(df.columns)[:10]}...")  # 显示前10列
        
        return df
        
    except Exception as e:
        print(f"❌ 加载 {file_path} 失败: {e}")
        return None

def filter_label_1(df: pd.DataFrame, dataset_name: str) -> pd.DataFrame:
    """
    过滤出label=1的数据
    """
    try:
        if 'label' not in df.columns:
            print(f"⚠️  {dataset_name} 中没有找到 'label' 列，可用列: {list(df.columns)}")
            print("   请检查列名，可能是 'Label', 'target', 'is_malicious' 等")
            return df
            
        original_count = len(df)
        filtered_df = df[df['label'] == 1].copy()
        filtered_count = len(filtered_df)
        
        print(f"📊 {dataset_name} 过滤结果:")
        print(f"   原始数据: {original_count} 条")
        print(f"   label=1: {filtered_count} 条")
        print(f"   保留比例: {filtered_count/original_count:.2%}")
        
        return filtered_df
        
    except Exception as e:
        print(f"❌ 过滤 {dataset_name} 失败: {e}")
        return df

def align_columns(dfs: List[pd.DataFrame], names: List[str]) -> List[pd.DataFrame]:
    """
    对齐所有数据框的列，处理列不匹配的情况
    """
    print("\n🔄 开始列对齐...")
    
    # 获取所有列的并集
    all_columns = set()
    for df in dfs:
        all_columns.update(df.columns)
    
    all_columns = sorted(list(all_columns))
    print(f"   总计发现 {len(all_columns)} 个不同的列")
    
    # 为每个数据框添加缺失的列
    aligned_dfs = []
    for i, (df, name) in enumerate(zip(dfs, names)):
        print(f"   处理 {name}...")
        
        # 添加缺失的列，默认值为0或NaN
        missing_cols = set(all_columns) - set(df.columns)
        if missing_cols:
            print(f"     添加 {len(missing_cols)} 个缺失列")
            for col in missing_cols:
                # 如果列名包含这些关键词，默认值为0，否则为NaN
                if any(keyword in col.lower() for keyword in ['count', 'ratio', 'score', 'rule_', 'malicious_']):
                    df[col] = 0
                else:
                    df[col] = np.nan
        
        # 重新排序列以保持一致性
        df_aligned = df[all_columns].copy()
        
        # 添加数据来源标识
        df_aligned['data_source'] = name
        
        aligned_dfs.append(df_aligned)
        print(f"     {name} 对齐后形状: {df_aligned.shape}")
    
    return aligned_dfs

def merge_datasets(features_path: str, maloss_path: str, robust_path: str, 
                  output_path: str = "merged_dataset.csv") -> bool:
    """
    主要的数据集合并函数
    """
    print("🚀 开始数据集合并流程...")
    print("="*60)
    
    # 1. 加载所有数据集
    print("\n📁 第1步: 加载数据集")
    features_df = load_and_validate_csv(features_path)
    maloss_df = load_and_validate_csv(maloss_path)
    robust_df = load_and_validate_csv(robust_path)
    
    # 检查是否都加载成功
    if features_df is None or maloss_df is None or robust_df is None:
        print("❌ 部分数据集加载失败，无法继续")
        return False
    
    # 2. 过滤maloss和robust数据集
    print("\n🔍 第2步: 数据过滤")
    print("   features.csv: 保留全部数据")
    maloss_filtered = filter_label_1(maloss_df, "maloss.csv")
    robust_filtered = filter_label_1(robust_df, "robust.csv")
    
    # 3. 列对齐
    dfs_to_merge = [features_df, maloss_filtered, robust_filtered]
    names = ["features", "maloss", "robust"]
    
    aligned_dfs = align_columns(dfs_to_merge, names)
    
    # 4. 合并数据
    print("\n🔗 第3步: 合并数据")
    try:
        merged_df = pd.concat(aligned_dfs, ignore_index=True, sort=False)
        print(f"✅ 合并成功!")
        print(f"   最终数据形状: {merged_df.shape}")
        
        # 统计各数据源的数据量
        print(f"\n📊 数据来源统计:")
        source_counts = merged_df['data_source'].value_counts()
        for source, count in source_counts.items():
            percentage = count / len(merged_df) * 100
            print(f"   {source}: {count:,} 条 ({percentage:.1f}%)")
            
    except Exception as e:
        print(f"❌ 数据合并失败: {e}")
        return False
    
    # 5. 保存结果
    print(f"\n💾 第4步: 保存到 {output_path}")
    try:
        merged_df.to_csv(output_path, index=False)
        print(f"✅ 数据已成功保存到: {output_path}")
        
        # 显示保存后的文件信息
        file_size = Path(output_path).stat().st_size / (1024 * 1024)  # MB
        print(f"   文件大小: {file_size:.2f} MB")
        
        return True
        
    except Exception as e:
        print(f"❌ 保存失败: {e}")
        return False

def show_data_preview(file_path: str, n_rows: int = 5):
    """
    显示合并后数据的预览
    """
    try:
        df = pd.read_csv(file_path)
        print(f"\n👀 数据预览 ({file_path}):")
        print("-" * 50)
        print(f"数据形状: {df.shape}")
        print(f"列数: {len(df.columns)}")
        
        # 显示前几行
        print(f"\n前 {n_rows} 行数据:")
        print(df.head(n_rows))
        
        # 显示数据类型统计
        print(f"\n数据类型统计:")
        dtype_counts = df.dtypes.value_counts()
        for dtype, count in dtype_counts.items():
            print(f"   {dtype}: {count} 列")
            
        # 显示缺失值统计
        missing_counts = df.isnull().sum()
        missing_cols = missing_counts[missing_counts > 0]
        if len(missing_cols) > 0:
            print(f"\n缺失值统计 (前10个):")
            for col, count in missing_cols.head(10).items():
                percentage = count / len(df) * 100
                print(f"   {col}: {count} ({percentage:.1f}%)")
        else:
            print(f"\n✅ 无缺失值")
            
    except Exception as e:
        print(f"❌ 预览失败: {e}")

def main():
    """
    主函数
    """
    # 配置文件路径
    features_path = "features.csv"
    maloss_path = "maloss.csv"  
    robust_path = "robust.csv"
    output_path = "merged_dataset.csv"
    
    print("🎯 NPM恶意包数据集合并工具")
    print("=" * 60)
    print("合并规则:")
    print("  • features.csv: 保留全部数据")
    print("  • maloss.csv:   仅保留 label=1 的数据")
    print("  • robust.csv:   仅保留 label=1 的数据")
    print("=" * 60)
    
    # 执行合并
    success = merge_datasets(features_path, maloss_path, robust_path, output_path)
    
    if success:
        print(f"\n🎉 合并完成! 结果保存在: {output_path}")
        
        # 显示预览
        show_data_preview(output_path)
        
        print(f"\n💡 使用建议:")
        print(f"   1. 检查合并后的数据质量")
        print(f"   2. 处理可能的缺失值")
        print(f"   3. 验证各数据源的特征是否对齐")
        print(f"   4. 如需要，可以删除 'data_source' 列")
        
    else:
        print(f"\n❌ 合并失败，请检查错误信息并重试")
        return 1
    
    return 0

if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)